import os
import pandas as pd

#数据提取
demand_train_A = 'data/demand_train_A.csv'
geo_topo = 'data/geo_topo.csv'
inventory_info_A = 'data/inventory_info_A.csv'
product_topo = 'data/product_topo.csv'
weight_A = 'data/weight_A.csv'

demand_train_A = pd.read_csv(demand_train_A)
geo_topo = pd.read_csv(geo_topo)
inventory_info_A = pd.read_csv(inventory_info_A)
product_topo = pd.read_csv(product_topo)
weight_A = pd.read_csv(weight_A)

demand_test_A = 'data/demand_test_A.csv'
demand_test_A = pd.read_csv(demand_test_A)
dfs = [demand_train_A,geo_topo,inventory_info_A,product_topo,weight_A,demand_test_A]

for df in dfs:
    if 'Unnamed: 0' in df.columns:
        df.drop(columns='Unnamed: 0',inplace=True)#删掉'Unnamed: 0'的列
    if 'ts' in df.columns:
        df = df.sort_values(by='ts')#按时间排序

#数据合并
all_data =pd.concat([demand_train_A,demand_test_A])
all_data = all_data.sort_values(by='ts')
all_data = all_data.reset_index().drop(columns='index')

#预测未来需求(这里直接用了参考值)
#真实值串14天
submission = demand_test_A

#fillna(method='ffill')为用前值填充空缺值
submission['yesterday_qty'] = submission.groupby('unit')['qty'].shift(1).fillna(method='ffill').reset_index().sort_index().set_index('index')
submission['diff_1'] = submission['qty'] - submission['yesterday_qty']#该产品当日使用量增量（即当日新增需求）
submission['qty'] = submission['diff_1']#用增量列代替使用量列

#shift_14列为当前产品当前日期十四天后的需求增量
submission['shift_14']=submission.groupby('unit')['qty'].shift(-14).fillna(0).reset_index().sort_index().set_index('index')
#生成十四天后需求增量列表
submission = submission[['unit','ts','shift_14']].rename(columns={'shift_14':'qty'}) 


#按照七天聚合
submission['dt'] = pd.to_datetime(submission['ts'])#时间列
submission['weekofyear'] = submission['dt'].dt.weekofyear#此列表示第几周
submission['year'] = submission['dt'].dt.year#年份列
submission_week = submission.copy()

#sum_qty为当前产品当前日期十四天后一周的需求增量之和（即本日后第15-21天需求增量之和）
submission_week = submission_week.groupby(['weekofyear','year','unit'],as_index=False).sum()
submission_week['sum_qty'] = submission_week['qty']
submission = pd.merge(submission_week,submission,on = ['weekofyear','year','unit'])


#提取每周一决策时需要的预测数据
submission['dayofweek'] = submission['dt'].dt.dayofweek
submission = submission[submission['dayofweek']==0]
submission = submission[['unit','ts','sum_qty']].rename(columns={'sum_qty':'qty'})


#根据未来需求消耗掉库存
init_inventory = inventory_info_A.set_index(['unit'])['qty'].to_dict()
def consume_init_inventory(arr,init_val):
    remain = init_val
    i = 0
    while remain>0 and i<len(arr):      #初始库存为0时直接return？？？库存有负值时跳出循环放弃补货？
        arr[i] = max(0,arr[i]-remain)   #决策日十四天消耗后可能还是不满足预测值需求？？
        remain -= arr[i]
        i+=1
        print(remain)
    return arr

r = []
for i,group in submission.groupby('unit'):#遍历每类产品做相应决策

    unit = group['unit'].values[0]
    init_val = init_inventory[unit]#当前产品初始库存水位
    
    group = group.sort_values(by='ts')
    qty_list = group['qty'].values#当前产品两周后需求增量（后15-21天需求增量之和）按日期排序
    # print(qty_list) 
    # print(init_val)
    qty_list = consume_init_inventory(qty_list,init_val)
    group['qty'] = qty_list
    r.append(group)

submission = pd.concat(r)  
print(submission)
